package com.cn.tous.auth.oauth2.token;

import com.nimbusds.jwt.JWTClaimsSet;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.oidc.IdTokenClaimNames;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.time.Instant;
import java.util.*;

public class OidcIdTokenGenerator {


    public static OidcIdToken generateIdToken(
            Object authenticatedPrincipal,
            String clientId,
            String nonce,
            OAuth2AccessToken accessToken, String issuer) {

        try {

            Instant now = accessToken.getIssuedAt();

            Instant expiresAt =  accessToken.getExpiresAt();

            // 构建 Claims
            JWTClaimsSet.Builder claimsBuilder = new JWTClaimsSet.Builder()
                    .issuer(issuer)
                    .subject(getUserSubject(authenticatedPrincipal)) // sub
                    .audience(Collections.singletonList(clientId))
                    .issueTime(Date.from(now))
                    .expirationTime(Date.from(expiresAt))
                    .notBeforeTime(Date.from(now))
                    .claim("auth_time", Date.from(now));

            // 添加 nonce（如果提供）
            if (nonce != null && !nonce.isEmpty()) {
                claimsBuilder.claim("nonce", nonce);
            }

//
//            // 添加用户信息（模拟）
//            claimsBuilder
//                    .claim("username", "张三")
//                    .claim("email", "zhangsan@example.com")
//                    .claim("role", "USER")
//                    .claim("department", "IT");

            // 如果有 access_token，添加 at_hash
            if (accessToken != null) {
                String atHash = generateAtHash(accessToken.getTokenValue());
                claimsBuilder.claim("at_hash", atHash);
            }

            JWTClaimsSet claims = claimsBuilder.build();



            // 构造 Spring Security 的 OidcIdToken 对象
            return OidcIdToken.withTokenValue(accessToken.getTokenValue())
                    .claim(IdTokenClaimNames.SUB, claims.getSubject())
                    .claim(IdTokenClaimNames.ISS, claims.getIssuer())
                    .claim(IdTokenClaimNames.AUD, claims.getAudience())
                    .claim(IdTokenClaimNames.EXP, claims.getExpirationTime().toInstant())
                    .claim(IdTokenClaimNames.IAT, claims.getIssueTime().toInstant())
                    .claim("nonce", nonce)
                    .build();

        } catch (Exception e) {
            throw new RuntimeException("Failed to generate ID Token", e);
        }
    }

    private static String getUserSubject(Object principal) {
        // 这里可以根据 UserDetailsService 返回的 UserDetails 提取 sub

        return principal.toString();
    }

    /**
     * 生成 at_hash（access token 的 SHA-256 哈希前128位）
     */
    private static String generateAtHash(String accessTokenValue) {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            byte[] hash = md.digest(accessTokenValue.getBytes(StandardCharsets.UTF_8));

            // 取前 128 位（16 字节）
            byte[] leftPart = Arrays.copyOfRange(hash, 0, 16);

            // Base64URL 编码
            return Base64.getUrlEncoder().withoutPadding().encodeToString(leftPart);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }



}
